/*
 * Copyright 2018- The Pixie Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include "src/carnot/exec/ml/eigen_test_utils.h"

Eigen::MatrixXf kmeans_test_data() {
  Eigen::MatrixXf points(60, 2);
  // Generated from following python code.
  // import numpy as np
  // cluster_centers = [np.array([0, -1]), np.array([-1, 1]), np.array([1, 1])]
  // points = np.zeros((60, 2))
  // for i, c in enumerate(cluster_centers):
  //   noise = np.random.normal(scale=0.1, size=(20, 2))
  //   points[i*20:(i+1)*20] = c + noise
  // for p in points:
  //   print(f'{p[0]}f,{p[1]}f,', end='')
  points << 0.14433300616089734f, -0.8793180153309f, 0.05949586562982306f, -0.8981571875566496f,
      0.045912439529482464f, -1.0450058889229272f, -0.06214011236839572f, -1.0729903835740358f,
      0.04603515003488658f, -0.8846887866931837f, -0.21980546539488605f, -1.0394569061388066f,
      -0.11780845308249231f, -0.9658251999323095f, -0.06926263481614962f, -0.9206451534476171f,
      0.22291289136062462f, -0.8592457137662223f, -0.10812720951639021f, -1.0892107762722554f,
      -0.1792534402223778f, -1.022743910533264f, 0.08399157815261842f, -0.9113827245847695f,
      0.02431466579983471f, -1.0468270478847939f, 0.08158150707715967f, -1.0924550521431924f,
      0.04362853314724771f, -0.9599307455959076f, 0.028563484501809924f, -0.8513829979744424f,
      -0.036427600113894364f, -1.1383261837884853f, -0.13833396207420565f, -1.1909609228754165f,
      -0.060383979420806914f, -0.9764700682474772f, -0.0069194228693483855f, -0.8361066508470854f,
      -0.783234939668269f, 1.1100489037514472f, -1.123889472411103f, 1.2161425654836795f,
      -1.0268450125054402f, 0.8491924139939968f, -0.9972586880166647f, 1.074351034836843f,
      -1.0905773300860053f, 0.856226352344647f, -1.0050295410202241f, 1.2255046702992871f,
      -0.9714789085720638f, 0.9142034339789499f, -1.1193719577190944f, 0.8552578791363454f,
      -0.9894167345565542f, 1.0715730469626137f, -0.8638485675777328f, 0.9747595411999916f,
      -1.2169100400390367f, 1.0294415716889667f, -1.063757220795692f, 0.9582668443125094f,
      -1.0208365714002077f, 0.9181076259280161f, -1.1268895780111996f, 1.0902284657376755f,
      -0.9326522338079375f, 1.1173884740891427f, -1.0990337208107066f, 1.097436043321699f,
      -1.0275829825694165f, 1.0837132671204637f, -1.0077841520188935f, 0.9217936949427348f,
      -1.076235551445322f, 1.1644627934094223f, -0.9240929922109229f, 1.0856715897935156f,
      0.9380426715036666f, 1.003694658627763f, 1.0302799680644827f, 0.962223201442501f,
      1.0080552686746307f, 0.94980432091717f, 1.074241365749615f, 1.0708860206145085f,
      0.8883968531271494f, 1.047886240539617f, 1.0124385021719202f, 0.9870085761270437f,
      0.97432159986717f, 0.9554950536442475f, 1.068325603686679f, 0.9037163434153809f,
      1.1571765385736936f, 1.0289673805333335f, 1.0501896473724097f, 1.0042518859922789f,
      0.8898872184254613f, 1.0044342709872092f, 0.9977986967044865f, 0.9571779620718711f,
      0.9813777282435779f, 1.0630417153018321f, 1.0852703236872203f, 1.0732736653082076f,
      1.051568934631669f, 1.0070787214519146f, 0.7988360565156847f, 1.0171809578331912f,
      1.0958762601607985f, 1.108771957412977f, 1.0513245047376318f, 1.1120995446221382f,
      0.912355238875605f, 1.0289852281261194f, 1.06201913161301f, 0.9294042130778405f;
  return points;
}

Eigen::MatrixXf kmeans_expected_centroids() {
  Eigen::MatrixXf expected_centroids(3, 2);
  expected_centroids << 0, -1, -1, 1, 1, 1;
  return expected_centroids;
}
